
import numpy as np
import math
import matplotlib.pyplot as plt


class CrossSection:
    def __init__(self, threshold_energy, reaction_name, is_ground_state = True):
        self.threshold_energy = threshold_energy

        data_temp = [[self.threshold_energy],[0.0]]
        for energy in np.arange(float(int(self.threshold_energy+1)), 100.0, 1.0):
            CS = 0.0
            data_temp[0].append(energy)
            data_temp[1].append(CS)

        for energy in np.arange(100.0, 1000.0, 10.0):
            CS = 0.0
            data_temp[0].append(energy)
            data_temp[1].append(CS)

        data_np = np.asarray(data_temp)
        data_np = np.transpose(data_np)

        #cs_data: cross section data
        self.cs_data = data_np.copy()
        self.collision_strength_data = data_np.copy()

        self.reaction_name = reaction_name

        self.is_ground_state = is_ground_state

    def get_cross_section(self, energy):
        cs_temp = 0.0
        if energy <= self.threshold_energy:
            return 0.0
        elif energy >= self.cs_data[-1,0]:
            return self.cs_data[-1,1]
        else:
            for i in range(2,self.cs_data.shape[0]-1):
                if self.cs_data[i+1,0] >= energy:
                    e0 = self.cs_data[i,0]
                    e1 = self.cs_data[i+1,0]
                    cs0 = self.cs_data[i,1]
                    cs1 = self.cs_data[i+1,1]
                    c0 = (e1-energy) / (e1-e0)
                    c1 = (energy-e0) / (e1-e0)
                    return cs0 * c0 + cs1 * c1
            print("Cross Section Error: get_cross_section")

class CS_Excitation:
    def __init__(self, species_name):
        self.Vif = 0.0
        self.min_Vif = 1000.0

        self.cs_list = []
        self.cs_total = None

        self.data_prefix = "data/"
        self.species_name = species_name




    def read_cross_section(self,reaction_name):
        filename_ori = "original/" + reaction_name + ".txt"
        data = np.loadtxt(filename_ori)
        data[:,1] = data[:,1] * 1.0e-20
        
        self.Vif = data[0,0]
        if self.Vif < self.min_Vif:
            self.min_Vif = self.Vif


        cs = CrossSection(self.Vif, reaction_name)
        cs.cs_data = data.copy()

        self.cs_list.append(cs)


    
              
    def plot_cs(self):
        for i in range(len(self.cs_list)):
            #if "Excitation_C_2s22p23P-2s2p33P" == self.cs_list[i].reaction_name:
            #    plt.semilogx(self.cs_list[i].cs_data[:,0], self.cs_list[i].cs_data[:,1])
            plt.plot(self.cs_list[i].cs_data[:,0], self.cs_list[i].cs_data[:,1])
            if self.cs_list[i].cs_data[:,1].min() < 0.0:
                print("cs below zero: ", self.cs_list[i].reaction_name, self.cs_list[i].cs_data[:,1].min())
            #for j in range(self.cs_list[i].cs_data[:,0].shape[0]):
                #print("cs energy: ", self.cs_list[i].cs_data[j,0], self.cs_list[i].collision_strength_data[j,0])
            #    print("cs energy: ", self.cs_list[i].cs_data[j,0], self.cs_list[i].cs_data[j,1])
        plt.xlim((0.0, 500))
        #plt.ylim((0.0, 10))
        #plt.legend()
        plt.savefig("fig/Excitation_" + self.species_name + ".svg")

    def save_cs(self):
        for i in range(len(self.cs_list)):
            np.savetxt(self.data_prefix + self.cs_list[i].reaction_name + "_nist.dat", self.cs_list[i].cs_data, fmt='%1.5e')


    def cal_total_cs(self):
        self.cs_total = CrossSection(self.min_Vif, "Excitation_" + self.species_name + "_total")
        for i in range(1, self.cs_total.cs_data.shape[0]):
            energy = self.cs_total.cs_data[i,0]
            cs_i = 0.0
            for j in range(len(self.cs_list)):
                cs_i = cs_i + self.cs_list[j].threshold_energy * self.cs_list[j].get_cross_section(energy)
            cs_i = cs_i / self.min_Vif
            self.cs_total.cs_data[i,1] = cs_i
        np.savetxt(self.data_prefix + self.cs_total.reaction_name + ".dat", self.cs_total.cs_data, fmt='%1.5e')

        plt.clf()
        plt.plot(self.cs_total.cs_data[:,0], self.cs_total.cs_data[:,1])
        plt.xlim((0.0, 1000))
        plt.savefig("fig/Excitation_" + self.species_name + "_total.svg")

if __name__ == '__main__':
    cs = CS_Excitation()
    cs.generate_cross_section()
    cs.plot_cs()
    cs.save_cs()
    cs.cal_total_cs()

